import joblib
import numpy as np
import time

from src.basic_model.word2vec_processor import Word2VecProcessor
from utils.config import Config

root_path = "E:/Python+AI/group4_nlp_project"
conf = Config(root_path)

# 模型路径
cat_model_path = conf.basic_model_path + "cat_rf.pkl"
label_model_path = conf.basic_model_path + "label_lr.pkl"

# 词向量模型路径
cat_w2v_path = conf.basic_model_path + "cat_word2vec.pkl"
label_w2v_path = conf.basic_model_path + "label_word2vec.pkl"

# 加载模型和词向量
rf_cat = joblib.load(cat_model_path)
lr_label = joblib.load(label_model_path)
w2v_cat = Word2VecProcessor.load(cat_w2v_path)
w2v_label = Word2VecProcessor.load(label_w2v_path)

# 🔹 类别映射表
cat_id2label = {
    0: "书籍", 1: "平板", 2: "手机", 3: "水果", 4: "洗发水",
    5: "热水器", 6: "蒙牛", 7: "衣服", 8: "计算机", 9: "酒店"
}
label_map = {0: "负面", 1: "正面"}


def basic_predict(texts: list, task: str = "all"):
    start_time = time.time()  # 记录开始时间
    results = {"items": [], "total": len(texts), "totalTime": 0}

    for text in texts:
        result = {}

        # 预测分类任务（category prediction）
        if task in ("cat", "all"):
            vec = w2v_cat.get_sentence_vector(str(text))
            if vec is None:
                result["cat"] = {"error": "文本无法生成向量"}
            else:
                vec = np.array(vec).reshape(1, -1)
                probs = rf_cat.predict_proba(vec)[0]  # 所有类别概率
                top_idx = np.argmax(probs)  # 取最大概率的类别
                best_category = cat_id2label.get(int(top_idx), str(top_idx))
                best_probability = float(probs[top_idx])
                result["cat"] = {
                    "best_category": best_category,
                    "best_probability": best_probability
                }

        # 预测情感任务（sentiment prediction）
        if task in ("label", "all"):
            vec = w2v_label.get_sentence_vector(str(text))
            if vec is None:
                result["label"] = {"error": "文本无法生成向量"}
            else:
                vec = np.array(vec).reshape(1, -1)
                probs = lr_label.predict_proba(vec)[0]
                pred = np.argmax(probs)
                sentiment = "positive" if pred == 1 else "negative"  # Assuming 1 is positive and 0 is negative
                result["label"] = {
                    "sentiment": sentiment,
                    "score": float(probs[pred]),
                    "tags": [result["cat"].get("best_category", "unknown")]  # Use best_category as the tag
                }

        # 计算推理时间 (ms)
        inference_time = round((time.time() - start_time) * 1000, 2)
        result["inferenceTime"] = inference_time
        results["totalTime"] += inference_time

        # Add the result for each text
        results["items"].append({
            "text": text,
            "sentiment": result["label"].get("sentiment", "neutral"),
            "score": round(result["label"].get("score", 0), 4),
            "tags": result["label"].get("tags", []),
            "inferenceTime": inference_time
        })

    # Return the final structure
    return {
        "modelId": "basic",  # 模型ID
        "modelName": "基线模型(逻辑回归+随机森林)",  # 模型名称
        "total": results["total"],  # 总预测条数
        "totalTime": results["totalTime"],  # 总耗时
        "items": results["items"]  # 单条预测结果列表
    }



if __name__ == "__main__":
    texts = [
        "屏幕显示很给力，系统优化也不错，很流畅，使用起来把苹果甩了几条街，虚拟键的滑动和点击返回功能非常好用，包装精致，京东自营第二天准时派件。唯一缺陷，指纹识别解锁功能偶尔会失灵，只能密码解锁。",
        "商品质量很好，物流也快，推荐购买！",
        "包装破损，体验很差"
    ]

    result = basic_predict(texts, task="all")
    print(result)

# {
#     'modelId': 'basic',
#     'modelName': '基线模型(逻辑回归+随机森林)',
#     'total': 3,
#     'totalTime': 2300.02,
#     'items':
#         [
#             {
#                 'text': '屏幕显示很给力，系统优化也不错，很流畅，使用起来把苹果甩了几条街，虚拟键的滑动和点击返回功能非常好用，包装精致，京东自营第二天准时派件。唯一缺陷，指纹识别解锁功能偶尔会失灵，只能密码解锁。',
#                 'sentiment': 'positive',
#                 'score': 0.8947,
#                 'tags': ['质量好', '物流快', '推荐'],
#                 'inferenceTime': 750.54
#             },
#             {
#                 'text': '商品质量很好，物流也快，推荐购买！',
#                 'sentiment': 'positive',
#                 'score': 0.9972,
#                 'tags': ['质量好', '物流快', '推荐'],
#                 'inferenceTime': 768.52
#             },
#             {
#                 'text': '包装破损，体验很差',
#                 'sentiment': 'negative',
#                 'score': 0.9693,
#                 'tags': ['质量好', '物流快', '推荐'],
#                 'inferenceTime': 780.96
#             }
#         ]
# }